package org.ricks.gateway.cluster.transport;

import org.ricks.ioc.utils.Logger;
import org.ricks.net.StateMachineEnum;
import org.ricks.net.extension.plugins.IdleStatePlugin;
import org.ricks.net.transport.AioQuickClient;
import org.ricks.net.transport.AioQuickServer;
import org.ricks.net.transport.AioSession;
import org.ricks.net.transport.WriteBuffer;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import io.scalecube.cluster.transport.api.Message;
import io.scalecube.cluster.transport.api.Transport;
import io.scalecube.cluster.transport.api.TransportConfig;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Map;
import java.util.Optional;
import java.util.Queue;
import java.util.UUID;
import java.util.concurrent.*;

import static io.scalecube.cluster.transport.api.Message.HEADER_CORRELATION_ID;

public class SmartSocketTransport implements Transport {
    private static final String HEADER_CORRELATION_ID = "correlation-id";

    private final TransportConfig config;
    private final AioQuickServer server;
    private final ConnectionManager connectionManager;
    private final Sinks.Many<Message> messageSink = Sinks.many().multicast().directBestEffort();
    private final ConcurrentHashMap<String, CompletableFuture<Message>> pendingRequests = new ConcurrentHashMap<>();
    private final Scheduler scheduler = Schedulers.boundedElastic();
    private volatile boolean stopped = true;

    public SmartSocketTransport(TransportConfig config) {
        this.config = config;
        this.server = new AioQuickServer(
                config.port(),
                new SmartSocketProtocol(),
                createServerProcessor()
        );
        this.connectionManager = new ConnectionManager(config);
    }

    private SmartSocketMessageProcessor createServerProcessor() {
        SmartSocketMessageProcessor processor = new SmartSocketMessageProcessor();
        processor.setListener(msg -> {
            String correlationId = msg.correlationId();
            if (correlationId != null && pendingRequests.containsKey(correlationId)) {
                CompletableFuture<Message> future = pendingRequests.remove(correlationId);
                future.complete(msg);
            } else {
                messageSink.tryEmitNext(msg);
            }
            return CompletableFuture.completedFuture(null);
        });
        return processor;
    }

    @Override
    public Mono<Transport> start() {
        return Mono.fromRunnable(() -> {
            if (!stopped) return;
            server.setReadBufferSize(4096);
            try {
                server.start();
            } catch (IOException e) {
                throw new TransportException("Failed to start server", e);
            }
            stopped = false;
        }).thenReturn(this);
    }

    @Override
    public Mono<Void> stop() {
        return Mono.fromRunnable(() -> {
            if (stopped) return;
            stopped = true;
            server.shutdown();
            connectionManager.close();
            messageSink.tryEmitComplete();
            pendingRequests.values().forEach(future ->
                    future.completeExceptionally(new TransportException("Transport stopped")));
            pendingRequests.clear();
        });
    }

    @Override
    public boolean isStopped() {
        return stopped;
    }

    @Override
    public String address() {
        return "0.0.0.0:" + config.port();
    }

    @Override
    public Mono<Void> send(String address, Message message) {
        return Mono.create(sink -> {
            try {
                ByteArrayOutputStream bos = new ByteArrayOutputStream();
                config.messageCodec().serialize(message, bos);
                byte[] serialized = bos.toByteArray();

                ByteBuffer buffer = ByteBuffer.allocate(4 + serialized.length);
                buffer.putInt(serialized.length);
                buffer.put(serialized);
                buffer.flip();

                connectionManager.getSession(address).thenAccept(session -> {
                    WriteBuffer writeBuffer = session.writeBuffer();
                    try {
                        writeBuffer.write(buffer.array());
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    writeBuffer.flush();
                }).exceptionally(ex -> {
                    sink.error(new TransportException("Connection failed", ex));
                    return null;
                });
            } catch (Exception e) {
                sink.error(new TransportException("Serialization failed", e));
            }
        });
    }

    @Override
    public Mono<Message> requestResponse(String address, Message message) {
        return Mono.create(sink -> {
            String correlationId = UUID.randomUUID().toString();
            Message request = Message.withData(message.data())
                    .qualifier(message.qualifier())
                    .correlationId(correlationId)
                    .build();

            CompletableFuture<Message> future = new CompletableFuture<>();
            pendingRequests.put(correlationId, future);

            // 设置超时
            Disposable timeout = scheduler.schedule(() -> {
                if (pendingRequests.remove(correlationId) != null) {
                    sink.error(new TimeoutException("Response timeout"));
                }
            }, 5000, TimeUnit.MILLISECONDS);

            future.whenComplete((response, error) -> {
                timeout.dispose();
                if (error != null) {
                    sink.error(error);
                } else {
                    sink.success(response);
                }
            });

            send(address, request).subscribe(
                    null,
                    ex -> {
                        pendingRequests.remove(correlationId);
                        sink.error(ex);
                    }
            );
        });
    }

    @Override
    public Flux<Message> listen() {
        return messageSink.asFlux();
    }

    // 优化后的连接管理
    private static class ConnectionManager {
        private final TransportConfig config;
        private final Map<String, CompletableFuture<AioSession>> sessionFutures = new ConcurrentHashMap<>();
        private final Map<String, Queue<AioSession>> sessionPools = new ConcurrentHashMap<>();
        private final Map<String, AioSession> activeSessions = new ConcurrentHashMap<>();

        public ConnectionManager(TransportConfig config) {
            this.config = config;
        }

        public CompletableFuture<AioSession> getSession(String address) {
            Queue<AioSession> pool = sessionPools.computeIfAbsent(address, k -> new ConcurrentLinkedQueue<>());
            AioSession session = pool.poll();

            if (session != null && !session.isInvalid()) {
                return CompletableFuture.completedFuture(session);
            }

            return sessionFutures.compute(address, (addr, future) -> {
                if (future != null && !future.isDone()) return future;

                CompletableFuture<AioSession> newFuture = new CompletableFuture<>();
                createSession(addr).whenComplete((sess, ex) -> {
                    if (ex != null) {
                        newFuture.completeExceptionally(ex);
                    } else {
                        activeSessions.put(addr, sess);
                        newFuture.complete(sess);
                    }
                });
                return newFuture;
            });
        }

        private CompletableFuture<AioSession> createSession(String address) {
            return CompletableFuture.supplyAsync(() -> {
                String[] parts = address.split(":");
                if (parts.length != 2) {
                    throw new TransportException("Invalid address: " + address);
                }

                String host = parts[0];
                int port = Integer.parseInt(parts[1]);

                CompletableFuture<AioSession> sessionFuture = new CompletableFuture<>();

                SmartSocketMessageProcessor processor = new SmartSocketMessageProcessor();
                processor.setSessionFuture(sessionFuture);
                processor.setListener(msg -> CompletableFuture.completedFuture(null));

                AioQuickClient client = new AioQuickClient(host, port, new SmartSocketProtocol(), processor);

                try {
                    client.start();
                    return sessionFuture.get(5, TimeUnit.SECONDS);
                } catch (InterruptedException | ExecutionException | TimeoutException | IOException e) {
                    throw new RuntimeException(e);
                }
            }).exceptionally(ex -> {
                throw new TransportException("Connection failed: " + address, ex);
            });
        }

        public void returnSession(String address, AioSession session) {
            if (session != null && !session.isInvalid()) {
                sessionPools.computeIfAbsent(address, k -> new ConcurrentLinkedQueue<>()).offer(session);
            }
        }

        public void invalidateSession(String address, AioSession session) {
            activeSessions.remove(address, session);
            if (session != null) {
                try {
                    session.close(false);
                } catch (Exception ignored) {}
            }
        }

        public void close() {
            activeSessions.values().forEach(session -> {
                try {
                    session.close(false);
                } catch (Exception ignored) {}
            });
            activeSessions.clear();
            sessionPools.clear();
            sessionFutures.clear();
        }
    }
}
